
import numpy as np
from PIL import Image


def roi(x,y, size=3):
    '''
    size is 3x3, x and y are center points
    '''
    yield (x,y) # Center 
    yield (x+1,y) # Right
    yield (x-1,y) # Left
    yield (x,y+1) # Above
    yield (x,y-1) # Below
    yield (x+1,y+1) # Above + Right 
    yield (x-1,y+1) # Above + Left 
    yield (x+1,y-1) # Below + Right
    yield (x-1,y-1) # Below + Left

def avg(i_img):
    '''
    run a 3x3 average across the image and return a new blurred image
    '''
    i_pix = i_img.load()
    newsize = (i_img.size[0], i_img.size[1])
    
    # Create a new image 
    o_img = Image.new(i_img.mode, newsize)
    o_pix = o_img.load()
    # Get the average and set pixel in new image
    for i in xrange(0,newsize[0]):
        for j in xrange(0,newsize[1]):
            sum = 0
            for x,y in roi(i,j):
                try:
                    sum += i_pix[x,y]
                except IndexError:
                    # one of the pixels is out of range of the image
                    # set to black to give us a border instead of resizing
                    sum = 0
            newpx = sum/9 # Average at center i,j
            o_pix[i,j] = newpx
            #o_img.putpixel((i,j),newpx)
    return o_img


def nearest_4nb(x,y):
    yield (x,y+1) # up
    yield (x+1,y) # right
    yield (x,y-1) # down
    yield (x-1,y) # left


def avg_mask(i_img):
    '''
    run a 3x3 spatial mask using the average of the four neighbors
    but excludes the center
    '''
    i_pix = i_img.load()
    newsize = (i_img.size[0], i_img.size[1])
    
    # Create a new image 
    o_img = Image.new(i_img.mode, newsize)
    o_pix = o_img.load()
    # Get the average and set pixel in new image
    for i in xrange(0,newsize[0]):
        for j in xrange(0,newsize[1]):
            sum = 0
            for x,y in nearest_4nb(i,j):
                try:
                    sum += i_pix[x,y]
                except IndexError:
                    # one of the pixels is out of range of the image
                    # set to black to give us a border instead of resizing
                    sum = 0
            newpx = sum/4 # Average of neighbors of i,j
            o_pix[i,j] = newpx
    return o_img


def Laplacian(i_pix, x,y):
    '''
    Returns the local result of the laplacian for a given pixel
    Index error occurs for pixels outside of the image
    '''
    return i_pix[x+1,y] + i_pix[x-1,y] + \
            i_pix[x,y+1] + i_pix[x,y-1] - \
            4*i_pix[x,y]

def Generate_Laplacian(i_img):
    i_size = i_img.size
    o_img = Image.new(i_img.mode, i_size)
    
    i_pix = i_img.load()
    o_pix = o_img.load()
    
    M,N = i_size
    for i in xrange(M):
        for j in xrange(N):
            try:
                o_pix[i,j] = Laplacian(i_pix, i, j)
            except IndexError:
                o_pix[i,j] = 0
    return o_img

# ToDo: Merge Laplacian image with original(?)


def get_hist(i_img):
    '''
    returns 'x' and 'bins' for plotting
    '''
    bins = i_img.histogram()
    x = [ i for i in range(len(bins))]
    return x, bins

def hist_eq(i_img):
    M,N = i_img.size
    MN = M*N
    
    # Create a copy 
    o_img = Image.new(i_img.mode, i_img.size)
    
    # Load a reference for pixel access
    i_pix = i_img.load()
    o_pix = o_img.load()

    # define the intensity levels
    L = 256
    
    # define the histogram distribution
    pr = []
    
    # get the bins/count for the intensity levels 
    bins = i_img.histogram()

    # find the normalized distribution
    for n in bins:
        pr.append(float(n)/float(MN))
    
    # find the cdf of the normalized distribution
    cdf = [] 
    for i in range(L):
        sumpr = 0.00
        for j in range(i):
            sumpr += pr[j]
        cdf.append(sumpr)

    # Compute the new output intensities for each pixel
    for x in range(M):
        for y in range(N):
            # get the current pixel value
            cv = i_pix[x,y]
            # set output pixel by scaling it with the cdf
            o_pix[x,y] = (L-1) * cdf[cv]
    return o_img





if __name__ == '__main__':
    from PIL import Image

    original = Image.open('../img/winter_sm.jpg')
    
    x = original.rotate(-90)
    
    gray = x.convert('L')
    
    ox,oy = get_hist(gray)
    heq = hist_eq(gray)
    hx,hy = get_hist(heq)
    arr_img1 = np.array(gray)
    arr_img2 = np.array(heq)
    

